package edu.protbioinfo;

import org.apache.commons.lang.StringUtils;
import org.apache.commons.lang.math.NumberUtils;

import edu.protbioinfo.SequenceAligner.Direction;
import edu.protbioinfo.SequenceAligner.ScoreMatrix.CellValue;
import edu.protbioinfo.SequenceAligner.ScoreMatrix.SourceDirection;
import edu.protbioinfo.Utils.Pair;

public class AffineGapSequenceAligner extends AbstractSequenceAligner {
	private ScoreMatrix matchMatrix;
	private ScoreMatrix gapXMatrix;
	private ScoreMatrix gapYMatrix;
	
	private double gapOpenPenality;
	private double gapExtendPenality;
	
	
	public double getGapOpenPenality() {
		return gapOpenPenality;
	}

	public void setGapOpenPenality(double gapOpenPenality) {
		this.gapOpenPenality = gapOpenPenality;
	}

	public double getGapExtendPenality() {
		return gapExtendPenality;
	}

	public void setGapExtendPenality(double gapExtendPenality) {
		this.gapExtendPenality = gapExtendPenality;
	}

	public ScoreMatrix getMatchMatrix() {
		return matchMatrix;
	}

	public ScoreMatrix getGapXMatrix() {
		return gapXMatrix;
	}

	public ScoreMatrix getGapYMatrix() {
		return gapYMatrix;
	}

	public SequenceAlignment align(Pair<String, String> source) {
		SequenceAlignment alignment = new SequenceAlignment();
		
		if(source != null && 
		   StringUtils.isNotEmpty(source.first()) && 
		   StringUtils.isNotEmpty(source.second())){
			
			matchMatrix = new ScoreMatrix(source);
			matchMatrix.setName("M");

			gapXMatrix = new ScoreMatrix(source);
			gapXMatrix.setName("Ix");
			
			gapYMatrix = new ScoreMatrix(source);
			gapYMatrix.setName("Iy");
						
			int m = matchMatrix.getRows();
			int n = matchMatrix.getCols();
			
			// construct scores for three matrices and find max score
			double maxScore = findMatchScore(m - 1, n - 1);
			ScoreMatrix startMatrix = matchMatrix;
			
			double curScore = findXGapScore(m - 1, n - 1);
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				startMatrix = gapXMatrix;
			}

			curScore = findYGapScore(m - 1, n - 1);
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				startMatrix = gapYMatrix;
			}
					
			alignment.setScore(maxScore);	
			
			traceback(startMatrix, alignment);
		}		
		return alignment;
	}

	private double findMatchScore(int i, int j) {
		double maxScore = NEGATIVE_INFINITY;
		SourceDirection sourceDirection = null;
		
		if(i == 0 && j == 0){
			maxScore = 0.;
			sourceDirection = new SourceDirection(matchMatrix, Direction.Stop);
		}
		else if(i == 0 || j == 0){
			maxScore = NEGATIVE_INFINITY;
			sourceDirection = new SourceDirection(matchMatrix, null);
		}
		else{
			
			double curScore = findMatchScore(i - 1, j - 1) + substitutionMatrix.getValue(matchMatrix.getRowChar(i), matchMatrix.getColChar(j));
			sourceDirection = new SourceDirection(matchMatrix, Direction.Diagonal);
			
			maxScore = curScore;
						
			curScore = findXGapScore(i - 1, j - 1) + substitutionMatrix.getValue(matchMatrix.getRowChar(i), matchMatrix.getColChar(j));
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				sourceDirection = new SourceDirection(gapXMatrix, Direction.Diagonal);
			}
			
			curScore = findYGapScore(i - 1, j - 1) + substitutionMatrix.getValue(matchMatrix.getRowChar(i), matchMatrix.getColChar(j));
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				sourceDirection = new SourceDirection(gapYMatrix, Direction.Diagonal);
			}		
		}
		
//		System.out.println("Set sourceDirect to matchMatrix=>" + sourceDirection);
		matchMatrix.setValue(i, j, new CellValue(sourceDirection, maxScore));
		return maxScore;
	}

	private double findXGapScore(int i, int j) {
		double maxScore = NEGATIVE_INFINITY;
		SourceDirection sourceDirection = null;

		if(i == 0 && j == 0){
			maxScore = 0.;
			sourceDirection = new SourceDirection(gapXMatrix, Direction.Stop);		
		}
		else if(i == 0){
			maxScore = NEGATIVE_INFINITY;
			sourceDirection = new SourceDirection(gapXMatrix, null);
		}
		else if(j == 0){
			maxScore = -1 * (gapOpenPenality + gapExtendPenality * (i - 1));
			sourceDirection = new SourceDirection(gapXMatrix, Direction.Up);
		}
		else{
			double curScore = findMatchScore(i - 1, j) - gapOpenPenality;
			sourceDirection = new SourceDirection(matchMatrix, Direction.Up);
			
			maxScore = curScore;
			
			curScore = findXGapScore(i - 1, j) - gapExtendPenality;
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				sourceDirection = new SourceDirection(gapXMatrix, Direction.Up);
			}
		}
		
//		System.out.println("Set sourceDirect to gapXMatrix=>" + sourceDirection);
		gapXMatrix.setValue(i, j, new CellValue(sourceDirection, maxScore));
		return maxScore;
	}

	private double findYGapScore(int i, int j) {
		double maxScore = NEGATIVE_INFINITY;
		SourceDirection sourceDirection = null;

		if(i == 0 && j == 0){
			maxScore = 0.;
			sourceDirection = new SourceDirection(gapYMatrix, Direction.Stop);		
		}
		else if(j == 0){
			maxScore = NEGATIVE_INFINITY;
			sourceDirection = new SourceDirection(gapYMatrix, null);
		}
		else if(i == 0){
			maxScore = -1 * (gapOpenPenality + gapExtendPenality * (j - 1));
			sourceDirection = new SourceDirection(gapYMatrix, Direction.Left);
		}
		else{
			double curScore = findMatchScore(i, j - 1) - gapOpenPenality;
			sourceDirection = new SourceDirection(matchMatrix, Direction.Left);

			maxScore = curScore;
			
			curScore = findYGapScore(i, j - 1) - gapExtendPenality;
			if(Double.compare(curScore, maxScore) > 0){
				maxScore = curScore;
				sourceDirection = new SourceDirection(gapYMatrix, Direction.Left);
			}
		}
		
		
//		System.out.println("Set sourceDirect to gapYMatrix=>" + sourceDirection);
		gapYMatrix.setValue(i, j, new CellValue(sourceDirection, maxScore));
		return maxScore;
	}

	private void traceback(ScoreMatrix startMatrix, SequenceAlignment alignment) {
		int i = startMatrix.getRows() - 1;
		int j = startMatrix.getCols() - 1;
		
		ScoreMatrix curMatrix = startMatrix;
		
		SourceDirection sourceDir = curMatrix.getValue(i, j).getSourceDirection();
	
		String firstSeq = "";
		String secondSeq = "";
	
		while(!Direction.Stop.equals(sourceDir.getDirection())){
			switch(sourceDir.getDirection()){
			case Diagonal:
				firstSeq  = curMatrix.getColChar(j) + firstSeq;
				secondSeq = curMatrix.getRowChar(i) + secondSeq;
	
				i--;
				j--;
				break;
			case Up:
				firstSeq  = SequenceAlignment.BLANK + firstSeq;
				secondSeq = curMatrix.getRowChar(i) + secondSeq;
				
				i--;
				break;
			case Left:
				firstSeq  = curMatrix.getColChar(j) + firstSeq;
				secondSeq = SequenceAlignment.BLANK + secondSeq;
				
				j--;										
				break;
			}
			
			curMatrix = sourceDir.getSource();
			sourceDir = curMatrix.getValue(i, j).getSourceDirection();
		}
		
		alignment.setAlignment(new Pair<String, String>(firstSeq, secondSeq));
	}

	public static void main(String[] args){
		AffineGapSequenceAligner aligner = new AffineGapSequenceAligner();		
		aligner.setSubstitutionMatrix(new LabeledMatrix<Double>(new char[]{'A', 'C', 'D', 'E'}, new char[]{'A', 'C', 'D', 'E'}){
			{
				this.setData(new Double[]{
						2., -2., -2., -1.,
						-2., 1., 0., 0.,
						-2., 0., 2., -2.,
						-1., 0., -2., 2.
				});
			}
			
		});
		
		aligner.setGapOpenPenality(1.8);
		aligner.setGapExtendPenality(0.4);
		
		SequenceAlignment seqAlign = aligner.align(new Pair<String, String>("AEECA", "CDAA"));
		
		
		System.out.println("Match Matrix:");
		System.out.println(aligner.getMatchMatrix().toPrettyString());

		System.out.println("Ix Matrix:");
		System.out.println(aligner.getGapXMatrix().toPrettyString());
		
		System.out.println("Iy Matrix:");
		System.out.println(aligner.getGapYMatrix().toPrettyString());
		
		System.out.println("Optimal Alignment:");
		System.out.println(seqAlign.toPrettyString());		
	}
}
